#!/usr/bin/env python

import os
import shutil
import subprocess
import time
import multiprocessing
import logging
import random
import glob
import signal
#Local imports
import loghelper
from config import ConfigValues
from config import split_strip
from filehelper import create_archive
from filehelper import move_results
from filehelper import store_dict
from filehelper import read_dict
from executionhelper import safe_exec
from executionhelper import chroot_mount
from executionhelper import safe_exec_fork
from executionhelper import remove_devices

# Split all parameters which have a list of options separated by ,
# For each option, a new parameter set is created and finally
# a set with all possible combinations is returned.
def permute_var(items):
	ret = []
	permuted = False
	for i in range(len(items)):
		if items[i][1].find(',') > -1:
			permuted = True
			for val in split_strip(',', items[i][1]):
				new_items = items[:]
				new_items[i] = (new_items[i][0],val)
				for stuff in permute_var(new_items):
					ret.append(stuff)
	if not permuted:
		ret.append(items)
	return ret

# replaces @NN@ with node id, @HH@ with node id in hex, @MAC@ with mac address and @MTU@ with MTU
def replace_string(s, nodeid, mac, mtu):	
	ret = s.replace('@NN@',str(nodeid))
	ret = ret.replace('@HH@',hex(nodeid))
	ret = ret.replace('@MAC@',mac)
	ret = ret.replace('@MTU@',mtu)

	if multiprocessing.cpu_count() > 1: # By default, schedule on all but last core
		cpu_core = "0-%s" % (multiprocessing.cpu_count()-2)
	else:
		cpu_core = "0"

	ret = ret.replace('@CPUS@',cpu_core)
	return ret

def calculate_run(config):
	ret = config.getint('general','repetitions')
	ret *= len(split_strip(',', config.get('general','experiments')))
	ret *= len(split_strip(',', config.get('general','topology')))
	bet = 0
	for i in split_strip(',', config.get('general','topology')):
		for j in permute_var(config.items(i)):
			bet += 1
	return ret*bet

def create_lxc_config(lxcconfig,nodeid,mac,mtu, config):
	output_name = lxcconfig
	newfile = open(output_name,'w')
	template = open(config.get('directories','configs')+'/lxc.conf','r')
	try:
		for s in template:
			if len(s.lstrip()) > 0 and s.lstrip()[0] == '#':
				continue
			s = replace_string(s, nodeid, mac, mtu)
			newfile.write(s)
	finally:
		newfile.close()
		template.close()

	return output_name

def create_devices(log, environment, config, mtu, total_nodes, use_chroot):
	try:
		if use_chroot:
			chroot_mount(log, config, environment)
	except:
		log.exception("Got error when mounting in chroot-containers")
		raise

	if not os.path.ismount("/sys/fs/cgroup"):
		log.info("Mounting /cgroup")
		if not os.path.isdir('/cgroup'):
			os.mkdir('/cgroup')
	
		if not os.path.ismount('/cgroup'):
			safe_exec(['mount','-t','cgroup','cgroup','/cgroup'], log, config, environment)
	
	log.info("creating devices with mtu %s", mtu)

	for n in range(1,total_nodes+1):
		tapdev='tap-'+str(n)
		brdev='br-'+str(n)
		node='node'+str(n)
		mac='00:16:3e:00:01:%0.2X' % int(n)

		log.info("Creating bridge %s and tap-device %s for node %s with mac %s", brdev,tapdev,node,mac)

		lxcconfig='/tmp/lxc-'+str(n)+'.conf'

		# set up virtual devices
		safe_exec(['brctl','addbr',brdev], log, config, environment)
		safe_exec(['tunctl','-t',tapdev], log, config, environment)
		safe_exec(['ifconfig',tapdev,'0.0.0.0','mtu',str(mtu),'promisc','up'], log, config, environment)
		safe_exec(['brctl','addif',brdev,tapdev], log, config, environment)
		safe_exec(['ifconfig',brdev,'up'], log, config, environment)
		safe_exec(['ifconfig',brdev,'mtu',mtu,'hw','ether',mac], log, config, environment)

		# write env
		environment[n]['tap_device'] = tapdev
		environment[n]['bridge_device'] = brdev
		environment[n]['node'] = node
		environment[n]['node_id'] = str(n)
		environment[n]['mac_address'] = mac

		# create lxc node
		create_lxc_config(lxcconfig,n,mac,mtu, config)
		safe_exec(['lxc-create','-n',node,'-f',lxcconfig], log, config, environment)

	for bnf in glob.glob('/proc/sys/net/bridge/bridge-nf-*'):
		log.info("echo 0 > %s", bnf)
		f = open(bnf,'w')
		f.write('0')
		f.close

# should be run outside lxc
def start_ns3(parameters, log, config, env, rand_seed):
	ns3_pid = None
	total_nodes = config.getint('general', 'total_nodes')
	duration = config.getint('general', 'experiment_duration')
	initial_wait = config.getint('general', 'initial_wait')
	shutdown_wait = config.getint('general', 'shutdown_wait')

	log.info("Starting ns3")

	ns3_path = os.path.abspath(config.get('ns3','path'))
	if not os.path.isdir(ns3_path) or not os.path.exists(ns3_path + '/waf'):
		raise Exception("ns3 not found in " + ns3_path + ". Change ns3.path in settings.cfg to point to your ns-3 installation.")

	scratch_path = ns3_path + '/scratch/'

	if not os.path.isdir(scratch_path):
		raise Exception("scratch directory not found in "+ns3_path)

	script = os.path.abspath(config.get('directories','topologies')) + '/' + parameters['script']

	if not os.path.exists(script):
		raise Exception("simulation script not found in " + script)

	shutil.copy(script,scratch_path)

	cwd = os.getcwd()
	os.chdir(ns3_path)
	try:
		safe_exec(['./waf','-d','optimized'],log, config, env, chdir=False) # compile first
	except:
		log.exception("ns3 compilation error. Exiting")
		raise

	# TODO remove extension properly, not just :-3
	command = ['./waf','-d','optimized','--run']
	s = "scratch/" + parameters['script'][:-3]
	for c in parameters.items():
		if c[0] != "script":		 
			s = s + ' --'+ c[0] + "=" + c[1]

	s = s + ' --nodes=%s' % total_nodes
	s = s + ' --seed=%s' % rand_seed 
	s = s + ' --duration=%s' % (duration+initial_wait+shutdown_wait)

	command.append(s)
	log.info("ns3 command: %s", command)


	ns3_pid = safe_exec_fork(command, log, config, env, '0', True,chdir=False)
	os.chdir(cwd)
	return ns3_pid

def save_env(config, use_chroot, nodeid, environment):
	if (int(nodeid) == 0): # node id 0 is core
		store_dict(config.get('directories','dumps')+"/lxc-environment.dat", environment)
	else:
		if use_chroot:
			store_dict(config.get('directories','chroot_working_dir')+"/lxc-environment.dat", environment)
		else:
			raise Exception("Warning: lxc container attempted to overwrite global lxc-environment.dat")

def read_env(config, nodeid, use_chroot):
	if (int(nodeid) == 0): # node id 0 is core
		environment = read_dict(config.get('directories','main')+"/lxc-environment.dat")
	else:
		if use_chroot:
			environment = read_dict(config.get('directories','chroot_working_dir')+"/lxc-environment.dat")
		else:
			environment = read_dict(config.get('directories','main')+"/lxc-environment.dat")
	return environment

def start_nodes(experiment, log, config, env):
	total_nodes = config.getint('general', 'total_nodes')
	use_chroot = config.getboolean('general','use_chroot')
	
	#NodeID should be 0 when this method is called
	save_env(config, use_chroot, '0', env)

	for n in range(1,total_nodes+1):
		log.info("starting lxc node %s", (n))
#		safe_exec_fork(['lxc-start','--name','node'+str(n),'./run.py','startnode',str(n),experiment])
		if use_chroot:
			script = config.get('directories','chroot_bindpoint') + "/startnode.py"
		else:
			script = "./startnode.py"
		subprocess.Popen(['lxc-execute','--name','node'+str(n),script, str(n), experiment, config.get('source_archiver', 'config_file')],close_fds=True)

# this should be run from within lxc
def start_experiment(experiment, log, config, environment):
	log.info("Starting experiment %s", experiment)

	log.debug("Setting environment variables for experiment %s", experiment)
	try:
		for k in config.items(experiment):
			environment[0][k[0]] = str(k[1])
	except Exception,err:
			log.warning('A problem occured when adding experiment variables to environment, all variables after this one will not be added, if your experiment is dependent upon thoses variables the experiment will fail!')
			log.exception(err)

	start_nodes(experiment, log, config, environment)

def stop_ns3(ns3_pid, log):
	if ns3_pid != None:
		log.info("Stopping ns3...")
		ns3_pid.terminate()
		if ns3_pid.is_alive:
			time.sleep(1)
		if ns3_pid.is_alive:
			os.kill(ns3_pid.pid,signal.SIGTERM)
			
def run_control_init(log, config, environment):
	log.info("Running controler init scripts")
	log.debug('Controler init scripts: %s', split_strip(',', config.get('general','control_init_scripts')))
	for script in split_strip(',', config.get('general','control_init_scripts')):
		log.info(script)
		if config.has_section(script):
			log.info('Settings has %s section, adding section to environment', script)
			for name, value in config.items(script):
				environment[0][name] = value
		safe_exec([config.get('directories','modules')+"/"+script], log, config, environment)

def run_experiment(config, use_chroot,  environment, initial_wait, duration, shutdown_wait, total_nodes):
	log = logging.getLogger('core.startup')
	
	loghelper.add_logfile(log, config.get('directories','dumps') + "/core_startup.log")
	
	log.info('Copying folder "main" to virtual machine')
	#safe_exec([config.get('directories', 'modules') + '/' + config.get('virtual_machine', 'file'), 'copy_folder', os.path.abspath(config.get('directories', 'main')) + '/.'], log, config, environment)

	# Archive configuration files
	create_archive(log, config, environment)
	
	log = logging.getLogger('core.experiment')
	# Create all parameter combinations for the experiments
	for rep in range(config.getint('general','repetitions')):
		rand_seed = random.randint(0,65535) # set new rand seed every repetition
		environment[0]['rand_seed'] = str(rand_seed) # Add random seed to env so init scripts can access
		for experiment in split_strip(',', config.get('general','experiments')):
			for topology in split_strip(',', config.get('general','topology')):
				for f in permute_var(config.items(topology)):
					
					safe_exec([config.get('directories', 'modules') + '/' + config.get('virtual_machine', 'file'), 'copy_folder', os.path.abspath(config.get('directories', 'main')) + '/.'], log, config, environment)

					loghelper.add_logfile(log, config.get('directories','dumps') + "/core.log")
					run_control_init(log, config, environment)

					# convert f to dict
					sim_params = {}
					for t in f:
						sim_params[t[0]] = t[1]

					log.info("%s,%s,%s,%s" % (rep,experiment,topology,sim_params))
					log.info("---Experiment parameters---")
					environment[0]['current_experiment'] = experiment
					log.info("experiment=%s" % experiment)
					environment[0]['current_repetition'] = str(rep)
					log.info("current_repetition=%s" % rep)
					environment[0]['current_topology'] = topology
					log.info("current_topology=%s" % topology)
					#environment[0]['current_rand_seed'] = str(rand_seed)
					log.info("rand_seed=%s" % rand_seed)

					state = random.getstate()
					random.seed(rand_seed)
					for n in range(1,total_nodes+1):
						environment[n]['current_rand_seed'] = str(random.randint(0,65535))
						environment[n]['node_id'] = str(n)
						log.info("node "+str(n)+" rand_seed=" + str(environment[n]['current_rand_seed']))
					random.setstate(state)

					log.info("---Topology parameters---")

					for kv in sim_params.items():
						log.info("%s=%s", kv[0], kv[1])
						environment[0][kv[0]] = kv[1]
						
					if config.has_section(experiment):
						log.debug('Adding experiment config to environment')
						for k in config.items(experiment):
							environment[0][k[0]] = str(k[1])
					
					save_env(config, False, '0', environment)
					safe_exec([config.get('directories', 'modules') + '/' + config.get('virtual_machine', 'file'), 'copy_file', os.path.abspath(config.get('directories','dumps')+"/lxc-environment.dat")], log, config, environment)
					
					try:
						#create_devices(log, environment, config, sim_params['mtu'], total_nodes, use_chroot)
						safe_exec([config.get('directories', 'modules') + '/' + config.get('virtual_machine', 'file'), 'create_devices'], log, config, environment, select_time=True)
					except:
						log.error('Unable to create virtual devices and linux containers. Try cleanup.')
						raise
					try:
						ns3_pid = start_ns3(sim_params, log, config, environment, rand_seed)
						#start_experiment(experiment, log, config, environment)
						safe_exec_fork([config.get('directories', 'modules') + '/' + config.get('virtual_machine', 'file'), 'start_nodes', 'startnode.py', experiment, config.get('source_archiver', 'config_file')], log, config, environment)
						time.sleep(initial_wait + duration + shutdown_wait)
					finally:
						stop_ns3(ns3_pid, log)
						#remove_devices(total_nodes, use_chroot, config, log, environment)
						safe_exec([config.get('directories', 'modules') + '/' + config.get('virtual_machine', 'file'), 'remove_devices'], log, config, environment)
						loghelper.flush_logfiles(log)
					
					safe_exec([config.get('directories', 'modules') + '/' + config.get('virtual_machine', 'file'), 'move_results', os.path.abspath(config.get('directories', 'dumps'))], log, config, environment)
					try:
						move_results(log, config, environment, experiment + "_" + str(rand_seed))
					except:
						log.error("unable to store results! aborting...")
						raise
